import zmq
import numpy as np
import cv2
import json
from multiprocessing import shared_memory
import subprocess
import atexit
import os
import sys
import time
import uuid
import platform
import traceback


def get_platform_shmem_name(name):
    """根据平台调整共享内存名称"""
    if platform.system() == "Windows":
        return name
    else:
        return name.lstrip("/")


class ImageProcessor:
    def __init__(self):
        self.rust_process = None
        self.ctx = zmq.Context()
        self.socket = self.ctx.socket(zmq.REQ)
        self.socket.RCVTIMEO = 10000  # 10秒超时
        # 获取共享内存
        self.shmem_cache = None  # 共享内存缓存
        self.shm = self._get_shared_memory(1024 * 1024 * 1024)
        # 确定IPC路径（跨平台）
        self.ipc_path = self._get_ipc_path()
        self._start_rust_process()
        atexit.register(self.cleanup)

    def _get_ipc_path(self):
        """获取适合当前操作系统的IPC路径"""
        if platform.system() == "Windows":
            return "tcp://127.0.0.1:5555"
        else:
            return "ipc:///tmp/image_processor.ipc"

    def _start_rust_process(self):
        """启动Rust处理进程"""
        # 获取Rust可执行文件路径
        current_dir = os.path.dirname(os.path.abspath(__file__))
        rust_dir = os.path.join(current_dir, "..", "rust_processor")
        target_dir = "debug"

        # 构建可执行文件名（跨平台）
        executable = "rust_processor"
        if platform.system() == "Windows":
            executable += ".exe"

        rust_bin = os.path.join(rust_dir, "target", target_dir, executable)

        # 如果未构建，则尝试构建
        
        print("Building Rust processor...")
        build_result = subprocess.run(
            ["cargo", "build"],
            cwd=rust_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if build_result.returncode != 0:
            raise RuntimeError(
                f"Failed to build Rust processor:\n"
                f"STDOUT: {build_result.stdout.decode()}\n"
                f"STDERR: {build_result.stderr.decode()}"
            )
        if not os.path.exists(rust_bin):
            raise RuntimeError(f"can not find {rust_bin}")
        # 启动Rust进程
        self.rust_process = subprocess.Popen(
            [rust_bin],
            cwd=rust_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        # 等待进程启动
        time.sleep(2)
        if self.rust_process.poll() is not None:
            stdout, stderr = self.rust_process.communicate()
            raise RuntimeError(
                f"Rust processor failed to start:\nSTDOUT: {stdout}\nSTDERR: {stderr}"
            )

        # 连接到IPC
        try:
            self.socket.connect(self.ipc_path)
            print(f"Connected to Rust processor at {self.ipc_path}")
        except zmq.ZMQError as e:
            raise RuntimeError(f"Failed to connect to IPC: {str(e)}")

    def _get_shared_memory(self, required_size: int) -> shared_memory.SharedMemory:
        """获取或创建共享内存块"""
        # 如果已有内存块且足够大，则复用
        if self.shmem_cache and self.shmem_cache.size >= required_size:
            return self.shmem_cache

        # 否则创建新内存块
        if self.shmem_cache:
            self._cleanup_shmem(self.shmem_cache)

        # 创建新的共享内存（大小为所需大小的两倍，最小1MB）
        size = max(required_size * 2, 1024 * 1024)
        shmem_id = f"pyimg"
        platform_shmem_id = get_platform_shmem_name(shmem_id)

        try:
            shm = shared_memory.SharedMemory(
                name=platform_shmem_id, create=True, size=size
            )
            self.shmem_cache = shm
            return shm
        except Exception as e:
            raise RuntimeError(f"Shared memory creation failed: {str(e)}")

    def process(self, command: str, image: np.ndarray) -> np.ndarray:
        """处理图像并返回结果"""
        # 准备图像数据
        if image.dtype != np.uint8:
            image = image.astype(np.uint8)

        # 编码为PNG
        success, buffer = cv2.imencode(".png", image)
        if not success:
            raise ValueError("Failed to encode image as PNG")

        buffer_bytes = buffer.tobytes()
        input_size = len(buffer_bytes)

        # 复制数据到共享内存
        try:
            self.shm.buf[:input_size] = buffer_bytes
        except Exception as e:
            raise RuntimeError(f"Failed to write to shared memory: {str(e)}")

        # 发送控制消息
        ctrl_msg = {
            "command": command,
            "input_size": input_size,
            "shmem_size": self.shm.size,
            "params": None,
        }

        try:
            self.socket.send_string(json.dumps(ctrl_msg))
        except zmq.ZMQError as e:
            raise RuntimeError(f"Failed to send control message: {str(e)}")

        # 接收响应
        try:
            response = self.socket.recv_string()
        except zmq.Again:
            raise TimeoutError("Processing timeout after 10 seconds")
        except Exception as e:
            raise RuntimeError(f"Communication error: {str(e)}")

        resp_data = json.loads(response)
        if resp_data["status"] != "success":
            error_msg = resp_data.get("error", "Unknown error")
            raise RuntimeError(f"Processing error: {error_msg}")

        # 从共享内存读取结果
        output_size = resp_data["output_size"]
        output_ptr = input_size  # 输出数据在输入数据之后
        output_data = bytes(self.shm.buf[output_ptr : output_ptr + output_size])

        # 解码图像
        result_img = cv2.imdecode(
            np.frombuffer(output_data, dtype=np.uint8), cv2.IMREAD_UNCHANGED
        )

        if result_img is None:
            # 尝试直接加载为数组
            height, width = image.shape[:2]
            result_img = np.frombuffer(output_data, dtype=np.uint8).reshape(
                (height, width, 3)
            )

        if result_img is None:
            raise RuntimeError("Failed to decode result image")

        return result_img

    def _cleanup_shmem(self, shm: shared_memory.SharedMemory):
        """安全清理共享内存"""
        try:
            shm.close()
            shm.unlink()
        except Exception as e:
            print(f"Warning: Failed to cleanup shared memory: {str(e)}")

    def cleanup(self):
        """清理资源"""
        try:
            self.socket.close()
            self.ctx.term()
        except:
            pass

        if self.shmem_cache:
            self._cleanup_shmem(self.shmem_cache)
            self.shmem_cache = None

        if self.rust_process and self.rust_process.poll() is None:
            try:
                self.rust_process.terminate()
                self.rust_process.wait(timeout=2)
            except:
                try:
                    self.rust_process.kill()
                except:
                    pass


if __name__ == "__main__":
    # 测试代码
    processor = ImageProcessor()

    # 创建测试图像
    test_img = np.zeros((480, 640, 3), dtype=np.uint8)
    cv2.putText(
        test_img, "Test Image", (50, 240), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 3
    )

    commands = ["grayscale", "invert", "blur", "rotate90", "canny_edge"]

    for cmd in commands:
        try:
            print(f"Processing {cmd}...")
            start_time = time.time()
            result = processor.process(cmd, test_img)
            elapsed = time.time() - start_time
            print(f"  Processed in {elapsed:.4f} seconds")

            # 保存结果
            cv2.imwrite(f"result_{cmd}.jpg", result)
        except Exception as e:
            print(f"Error processing {cmd}: {str(e)}")
            traceback.print_exc()

    print("All processing completed. Check result images.")
